package com.gitee.wsl.platform.thread.task.callback

import com.gitee.wsl.platform.exception.logOrRethrow
import com.gitee.wsl.platform.exception.rethrowIfFatal
import com.gitee.wsl.platform.thread.task.TaskExecutor
import com.gitee.wsl.platform.thread.task.TaskLocalContext.signalTheStartOfBlockingCall
import com.gitee.wsl.platform.concurrent.api.Cancellable
import com.gitee.wsl.platform.thread.task.bean.Outcome
import com.gitee.wsl.platform.thread.task.collection.ImmutableStack
import org.jetbrains.annotations.ApiStatus
import java.io.Serializable
import java.time.Duration
import java.util.*
import java.util.concurrent.ExecutionException
import java.util.concurrent.TimeoutException
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicReference
import java.util.concurrent.locks.AbstractQueuedSynchronizer


/**
 * Represents a callback that will be invoked when a task completes.
 *
 *
 * A task can complete either successfully with a value, or with an exception,
 * or it can be cancelled.
 *
 *
 * MUST BE idempotent AND thread-safe.
 *
 * @param <T> is the type of the value that the task will complete with
</T> */
@FunctionalInterface
fun interface CompletionCallback<T> : Serializable {
    fun onOutcome(outcome: Outcome<T>)

    /**
     * Must be called when the task completes successfully.
     *
     * @param value is the successful result of the task, to be signaled
     */
    fun onSuccess(value: T) {
        onOutcome(Outcome.success(value))
    }

    /**
     * Must be called when the task completes with an exception.
     *
     * @param e is the exception that the task failed with
     */
    fun onFailure(e: Throwable) {
        onOutcome(Outcome.failure(e))
    }

    /**
     * Must be called when the task is cancelled.
     */
    fun onCancellation() {
        onOutcome(Outcome.cancellation())
    }

    companion object {
        /**
         * Returns a `CompletionListener` that does nothing (a no-op).
         */
        fun <T> empty(): CompletionCallback<T> {
            return CompletionCallback { outcome: Outcome<T> ->
                if (outcome is Outcome.Failure<T>) {
                    outcome.exception.logOrRethrow()
                }
            }
        }
    }
}

@ApiStatus.Internal
internal class ManyCompletionCallback<T> : CompletionCallback<T> {
    private val listeners: ImmutableStack<CompletionCallback<T>>

    @SafeVarargs
    constructor(vararg listeners: CompletionCallback<T>) {
        var stack: ImmutableStack<CompletionCallback<T>> = ImmutableStack.empty()
        for (listener in listeners) {
            stack = stack.prepend(listener)
        }
        this.listeners = stack
    }

    private constructor(
        listeners: ImmutableStack<CompletionCallback<T>>,
    ) {
        this.listeners = listeners
    }

    fun withExtraListener(extraListener: CompletionCallback<T>): ManyCompletionCallback<T> {
        val newListeners = this.listeners.prepend(extraListener)
        return ManyCompletionCallback(newListeners)
    }

    override fun onOutcome(outcome: Outcome<T>) {
        for (listener in listeners) {
            try {
                listener.onOutcome(outcome)
            } catch (e: Throwable) {
                e.logOrRethrow()
            }
        }
    }

    override fun onSuccess(value: T) {
        for (listener in listeners) {
            try {
                listener.onSuccess(value)
            } catch (e: Throwable) {
                e.logOrRethrow()
            }
        }
    }

    override fun onFailure(e: Throwable) {
        for (listener in listeners) {
            try {
                listener.onFailure(e)
            } catch (e: Throwable) {
                e.logOrRethrow()
            }
        }
    }

    override fun onCancellation() {
        for (listener in listeners) {
            try {
                listener.onCancellation()
            } catch (e: Throwable) {
                e.logOrRethrow()
            }
        }
    }
}

@ApiStatus.Internal
internal interface ContinuationCallback<T> : CompletionCallback<T>, Serializable {
    /**
     * Registers an extra callback to be invoked when the task completes.
     * This is useful for chaining callbacks or adding additional listeners.
     */
    fun registerExtraCallback(extraCallback: CompletionCallback<T>)
}

@ApiStatus.Internal
internal class AsyncContinuationCallback<T>(
    listener: CompletionCallback<T>,
    private val executor: TaskExecutor,
) : ContinuationCallback<T>, Runnable {
    private val isWaiting = AtomicBoolean(true)
    private val listenerRef: AtomicReference<CompletionCallback<T>> = AtomicReference<CompletionCallback<T>>(
        Objects.requireNonNull<CompletionCallback<T>>(listener, "listener")
    )

    private var outcome: Outcome<T>? = null
    private var successValue: T? = null
    private var failureCause: Throwable? = null
    private var isCancelled = false

    override fun run() {
        if (this.outcome != null) {
            listenerRef.get()!!.onOutcome(this.outcome!!)
        } else if (this.failureCause != null) {
            listenerRef.get()!!.onFailure(this.failureCause!!)
        } else if (this.isCancelled) {
            listenerRef.get()!!.onCancellation()
        } else {
            listenerRef.get()!!.onSuccess(this.successValue!!)
        }
        // For GC purposes; but it doesn't really matter if we nullify these or not
        this.outcome = null
        this.successValue = null
        this.failureCause = null
        this.isCancelled = false
    }

    override fun onOutcome(outcome: Outcome<T>) {
        Objects.requireNonNull<Any?>(outcome, "outcome")
        if (isWaiting.getAndSet(false)) {
            this.outcome = outcome
            executor.resumeOnExecutor(this)
        } else if (outcome is Outcome.Failure<*>) {
            outcome.exception.logOrRethrow()
        }
    }

    override fun onSuccess(value: T) {
        if (isWaiting.getAndSet(false)) {
            this.successValue = value
            executor.resumeOnExecutor(this)
        }
    }

    override fun onFailure(e: Throwable) {
        Objects.requireNonNull<Throwable>(e, "e")
        if (isWaiting.getAndSet(false)) {
            this.failureCause = e
            executor.resumeOnExecutor(this)
        } else {
            e.logOrRethrow()
        }
    }

    override fun onCancellation() {
        if (isWaiting.getAndSet(false)) {
            this.isCancelled = true
            executor.resumeOnExecutor(this)
        }
    }

    override fun registerExtraCallback(extraCallback: CompletionCallback<T>) {
        while (true) {
            val current = listenerRef.get()
            if (current is ManyCompletionCallback<T>) {
                val update = current.withExtraListener(extraCallback)
                if (listenerRef.compareAndSet(current, update)) {
                    return
                }
            } else if (listenerRef.compareAndSet(current, ManyCompletionCallback<T>(current, extraCallback))) {
                return
            }
        }
    }

    companion object {
        fun <T> protect(
            executor: TaskExecutor,
            listener: CompletionCallback<T>,
        ): ContinuationCallback<T> {
            return AsyncContinuationCallback<T>(
                listener,
                executor
            )
        }
    }
}

/**
 * INTERNAL API.
 *
 *
 * **INTERNAL API:** Internal apis are subject to change or removal
 * without any notice. When code depends on internal APIs, it is subject to
 * breakage between minor version updates.
 */
@ApiStatus.Internal
internal class BlockingCompletionCallback<T> : AbstractQueuedSynchronizer(), ContinuationCallback<T> {
    private val isDone = AtomicBoolean(false)
    private val extraCallbackRef = AtomicReference<CompletionCallback<T>?>(null)

    private var result: T? = null
    private var error: Throwable? = null
    private var interrupted: InterruptedException? = null

    private fun notifyOutcome() {
        val extraCallback = extraCallbackRef.getAndSet(null)
        if (extraCallback != null) try {
            if (error != null) extraCallback.onFailure(error!!)
            else if (interrupted != null) extraCallback.onCancellation()
            else extraCallback.onSuccess(result!!)
        } catch (e: Throwable) {
            e.logOrRethrow()
        }
        releaseShared(1)
    }

    override fun onSuccess(value: T) {
        if (!isDone.getAndSet(true)) {
            result = value
            notifyOutcome()
        }
    }

    override fun onFailure(e: Throwable) {
        e.rethrowIfFatal()
        if (!isDone.getAndSet(true)) {
            error = e
            notifyOutcome()
        } else {
            e.logOrRethrow()
        }
    }

    override fun onCancellation() {
        if (!isDone.getAndSet(true)) {
            interrupted = InterruptedException("Task was cancelled")
            notifyOutcome()
        }
    }

    override fun onOutcome(outcome: Outcome<T>) {
        if (outcome is Outcome.Success<T>) {
            onSuccess(outcome.value)
        } else if (outcome is Outcome.Failure<T>) {
            onFailure(outcome.exception!!)
        } else {
            onCancellation()
        }
    }

    override fun tryAcquireShared(arg: Int): Int {
        return if (state != 0) 1 else -1
    }

    override fun tryReleaseShared(arg: Int): Boolean {
        state = 1
        return true
    }

    internal fun interface AwaitFunction {
        @Throws(InterruptedException::class, TimeoutException::class)
        fun apply(isCancelled: Boolean)
    }

    @Throws(InterruptedException::class, ExecutionException::class, TimeoutException::class)
    private fun awaitInline(cancelToken: Cancellable, await: AwaitFunction): T {
        signalTheStartOfBlockingCall()
        var isCancelled = false
        var timedOut: TimeoutException? = null
        while (true) {
            try {
                await.apply(isCancelled)
                break
            } catch (e: TimeoutException) {
                if (!isCancelled) {
                    isCancelled = true
                    timedOut = e
                    cancelToken.cancel()
                }
            } catch (e: InterruptedException) {
                if (!isCancelled) {
                    isCancelled = true
                    cancelToken.cancel()
                }
            }
            // Clearing the interrupted flag may not be necessary,
            // but doesn't hurt, and we should have a cleared flag before
            // re-throwing the exception
            //
            // noinspection ResultOfMethodCallIgnored
            Thread.interrupted()
        }
        if (timedOut != null) throw timedOut
        if (interrupted != null) throw interrupted!!
        if (error != null) throw ExecutionException(error)
        return result!!
    }

    @Throws(InterruptedException::class, ExecutionException::class)
    fun await(cancelToken: Cancellable): T {
        try {
            return awaitInline(cancelToken, AwaitFunction { isCancelled: Boolean -> acquireSharedInterruptibly(1) })
        } catch (e: TimeoutException) {
            throw IllegalStateException("Unexpected timeout", e)
        }
    }

    @Throws(ExecutionException::class, InterruptedException::class, TimeoutException::class)
    fun await(cancelToken: Cancellable, timeout: Duration): T {
        return awaitInline(cancelToken, AwaitFunction { isCancelled: Boolean ->
            if (!isCancelled) {
                if (!tryAcquireSharedNanos(1, timeout.toNanos())) {
                    throw TimeoutException("Task timed-out after $timeout")
                }
            } else {
                // Waiting without a timeout, since at this point it's waiting
                // on the cancelled task to finish
                acquireSharedInterruptibly(1)
            }
        })
    }

    override fun registerExtraCallback(extraCallback: CompletionCallback<T>) {
        while (true) {
            val current = extraCallbackRef.get()
            if (current == null) {
                if (extraCallbackRef.compareAndSet(null, extraCallback)) {
                    return
                }
            } else if (current is ManyCompletionCallback<T>) {
                val update = current.withExtraListener(extraCallback)
                if (extraCallbackRef.compareAndSet(current, update)) {
                    return
                }
            } else if (extraCallbackRef.compareAndSet(current, ManyCompletionCallback<T>(current, extraCallback))) {
                return
            }
        }
    }
}
